package com.example.filter;

import com.example.config.EncryptionConfig;
import com.example.util.AesUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.util.ContentCachingResponseWrapper;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

public class ResponseEncryptFilter implements Filter {

    private static final Logger logger = LoggerFactory.getLogger(ResponseEncryptFilter.class);

    private final EncryptionConfig encryptionConfig;

    // 构造函数注入
    public ResponseEncryptFilter(EncryptionConfig encryptionConfig) {
        this.encryptionConfig = encryptionConfig;
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;

        // 检查请求头是否指示需要加密响应
        if ("true".equalsIgnoreCase(httpRequest.getHeader("X-Encrypt-Response")) && encryptionConfig.isEncryptionEnabled()) {
            ContentCachingResponseWrapper responseWrapper = new ContentCachingResponseWrapper(httpResponse);

            chain.doFilter(request, responseWrapper);

            String originalResponseBody = new String(responseWrapper.getContentAsByteArray(), httpResponse.getCharacterEncoding());

            if (!originalResponseBody.isEmpty()) {
                String encryptedBodyResponse = AesUtil.encrypt(encryptionConfig.getKey(), encryptionConfig.getInitVector(), originalResponseBody);
                if (encryptedBodyResponse != null) {
                    responseWrapper.resetBuffer();
                    responseWrapper.getWriter().write(encryptedBodyResponse);
                    responseWrapper.copyBodyToResponse();
                }
                logger.debug("加密后的响应体: {}", encryptedBodyResponse);
            }
        } else {
            // 如果不需要加密，继续处理原始请求
            chain.doFilter(request, response);
        }
    }
}

